/** @file fftw_impl.h
 * @brief Implementation of fftw-enabled FFT engines
 *
 */
#ifndef QMCPLUSPLUS_FFTW_IMPL_H
#define QMCPLUSPLUS_FFTW_IMPL_H
#include <fftw3.h>

/** fftw_base to be specialized with the precision
 */
template<typename T> struct fftw_base { };

template<>
struct fftw_base<double>
{
  typedef fftw_plan fft_plan_type;
  fft_plan_type forward_plan;
  fft_plan_type backward_plan;

  ///default constructor
  fftw_base() :forward_plan(0), backward_plan(0) { }

  ///virtual destructor to clean up the plans
  virtual ~fftw_base() { clean(); }

  ///clean the plans
  void clean()
  {
    if(forward_plan) {fftw_destroy_plan(forward_plan); forward_plan=0;}
    if(backward_plan) {fftw_destroy_plan(backward_plan); backward_plan=0;}
  }

  /** create complex-to-complex plan
   * @param dims fft size
   * @param howmany number of concurrent ffts
   * @param in complex input data
   * @param out complex output data
   * @param dir forward or backward direction
   * @param uflag fftw plan (measure ...)
   */
  fft_plan_type create_plan(int dims, int howmany
      , fftw_complex* in, fftw_complex* out , int dir, unsigned uflag)
  {
    if(howmany>1)
    {
      fftw_iodim data_dims;
      fftw_iodim howmany_dims;
      data_dims.n=dims; data_dims.is=1; data_dims.os=1;
      howmany_dims.n=howmany; howmany_dims.is=dims; howmany_dims.os=dims;
      return fftw_plan_guru_dft(1,&data_dims, 1,&howmany_dims, in, out,dir,uflag);
    }
    else
      return fftw_plan_dft_1d(dims,in,out,dir,uflag);
  }

  /** create real-to-complex forward plan 
   * @param dims fft size
   * @param howmany number of concurrent ffts
   * @param in real input data
   * @param out complex output data
   * @param uflag fftw plan (measure ...)
   */
  fft_plan_type create_plan(int dims, int howmany, double* in, fftw_complex* out, unsigned uflag)
  {
    if(howmany>1)
    {
      fftw_iodim data_dims;
      fftw_iodim howmany_dims;
      data_dims.n=dims; data_dims.is=1; data_dims.os=1;
      howmany_dims.n=howmany; howmany_dims.is=dims+2; howmany_dims.os=dims/2+1;
      return fftw_plan_guru_dft_r2c(1,&data_dims, 1,&howmany_dims , in, out,uflag);
    }
    else
      return fftw_plan_dft_r2c_1d(dims,in,out,uflag);
  }

  /** create complex-to-real backward plan 
   * @param dims fft size
   * @param howmany number of concurrent ffts
   * @param in complex input data
   * @param out real output data
   * @param uflag fftw plan (measure ...)
   */
  fft_plan_type create_plan(int dims, int howmany, fftw_complex* in, double* out, unsigned uflag)
  {
    if(howmany>1)
    {
      fftw_iodim data_dims;
      fftw_iodim howmany_dims;
      data_dims.n=dims; data_dims.is=1; data_dims.os=1;
      howmany_dims.n=howmany; howmany_dims.is=dims/2+1; howmany_dims.os=dims+2;
      return fftw_plan_guru_dft_c2r(1,&data_dims, 1,&howmany_dims,in,out,uflag);
    }
    else
      return fftw_plan_dft_c2r_1d(dims,in,out,uflag);
  }

  inline void execute_dft(fft_plan_type p, fftw_complex* in, fftw_complex* out)
  {
    fftw_execute_dft(p,in,out);
  }

  inline void execute_dft(fft_plan_type p, double* in, fftw_complex* out)
  {
    fftw_execute_dft_r2c(p,in,out);
  }

  inline void execute_dft(fft_plan_type p, fftw_complex* in, double* out)
  {
    fftw_execute_dft_c2r(p,in,out);
  }

  template<typename T1>
  inline fftw_complex* mangle(T1* in)
  {
    return reinterpret_cast<fftw_complex*>(in);
  }
};

template<>
struct fftw_base<float>
{
  typedef fftwf_plan fft_plan_type;
  fft_plan_type forward_plan;
  fft_plan_type backward_plan;

  ///default constructor
  fftw_base() :forward_plan(0), backward_plan(0) { }

  ///virtual destructor to clean up the plans
  virtual ~fftw_base() { clean(); }

  ///clean the plans
  void clean()
  {
    if(forward_plan) {fftwf_destroy_plan(forward_plan); forward_plan=0;}
    if(backward_plan) {fftwf_destroy_plan(backward_plan); backward_plan=0;}
  }

  /** create complex-to-complex plan
   * @param dims fft size
   * @param howmany number of concurrent ffts
   * @param in complex input data
   * @param out complex output data
   * @param dir forward or backward direction
   * @param uflag fftw plan (measure ...)
   */
  fft_plan_type create_plan(int dims, int howmany
      , fftwf_complex* in, fftwf_complex* out , int dir, unsigned uflag)
  {
    if(howmany>1)
    {
      fftw_iodim data_dims;
      fftw_iodim howmany_dims;
      data_dims.n=dims; data_dims.is=1; data_dims.os=1;
      howmany_dims.n=howmany; howmany_dims.is=dims; howmany_dims.os=dims;
      return fftwf_plan_guru_dft(1,&data_dims, 1,&howmany_dims, in, out,dir,uflag);
    }
    else
      return fftwf_plan_dft_1d(dims,in,out,dir,uflag);
  }

  /** create real-to-complex forward plan 
   * @param dims fft size
   * @param howmany number of concurrent ffts
   * @param in real input data
   * @param out complex output data
   * @param uflag fftw plan (measure ...)
   */
  fft_plan_type create_plan(int dims, int howmany, float* in, fftwf_complex* out, unsigned uflag)
  {
    if(howmany>1)
    {
      fftw_iodim data_dims;
      fftw_iodim howmany_dims;
      data_dims.n=dims; data_dims.is=1; data_dims.os=1;
      howmany_dims.n=howmany; howmany_dims.is=dims+2; howmany_dims.os=dims/2+1;
      return fftwf_plan_guru_dft_r2c(1,&data_dims, 1,&howmany_dims , in, out,uflag);
    }
    else
      return fftwf_plan_dft_r2c_1d(dims,in,out,uflag);
  }

  /** create complex-to-real backward plan 
   * @param dims fft size
   * @param howmany number of concurrent ffts
   * @param in complex input data
   * @param out real output data
   * @param uflag fftw plan (measure ...)
   */
  fft_plan_type create_plan(int dims, int howmany, fftwf_complex* in, float* out, unsigned uflag)
  {
    if(howmany>1)
    {
      fftw_iodim data_dims;
      fftw_iodim howmany_dims;
      data_dims.n=dims; data_dims.is=1; data_dims.os=1;
      howmany_dims.n=howmany; howmany_dims.is=dims/2+1; howmany_dims.os=dims+2;
      return fftwf_plan_guru_dft_c2r(1,&data_dims, 1,&howmany_dims,in,out,uflag);
    }
    else
      return fftwf_plan_dft_c2r_1d(dims,in,out,uflag);
  }

  inline void execute_dft(fft_plan_type p, fftwf_complex* in, fftwf_complex* out)
  {
    fftwf_execute_dft(p,in,out);
  }

  inline void execute_dft(fft_plan_type p, float* in, fftwf_complex* out)
  {
    fftwf_execute_dft_r2c(p,in,out);
  }

  inline void execute_dft(fft_plan_type p, fftwf_complex* in, float* out)
  {
    fftwf_execute_dft_c2r(p,in,out);
  }

  template<typename T1>
  inline fftwf_complex* mangle(T1* in)
  {
    return reinterpret_cast<fftwf_complex*>(in);
  }
};

/** specialization for single-fft only with FFTW on complex-to-complex
 */
template<typename T>
struct FFT_DFT<T,std::complex<T>,false,FFTW_ENG> 
: public fft_base, public fftw_base<T>
{
  typedef std::complex<T> value_type;
  using fftw_base<T>::forward_plan;
  using fftw_base<T>::backward_plan;

  FFT_DFT(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    init(dims,in,1,uflag);
  }

  void init(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    if(dims != fft_size)
    {
      fft_base::resize(dims,1);
#pragma omp critical
      {
        fftw_base<T>::clean();
        forward_plan = fftw_base<T>::create_plan(dims,1,mangle(in),mangle(in),FFTW_FORWARD,uflag);
        backward_plan = fftw_base<T>::create_plan(dims,1,mangle(in),mangle(in),FFTW_BACKWARD,uflag);
      }
    }
  }

  inline void transform_forward(value_type* in, int howmany=1)
  {
    value_type* offset=in;
    for(int i=0; i<howmany; ++i,offset+=fft_size) 
      execute_dft(forward_plan,mangle(offset),mangle(offset));
  }

  inline void transform_backward(value_type* in, int howmany=1)
  {
    value_type* offset=in;
    for(int i=0; i<howmany; ++i,offset+=fft_size) 
      execute_dft(backward_plan,mangle(offset),mangle(offset));
  }
};

/** specialization for single-fft only with FFTW on double
 */
template<typename T>
struct FFT_DFT<T,T,false,FFTW_ENG> 
: public fft_base, public fftw_base<T>
{
  typedef T value_type;
  using fftw_base<T>::forward_plan;
  using fftw_base<T>::backward_plan;
  int fft_mem_offset;

  FFT_DFT(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    init(dims,in,1,uflag);
  }

  void init(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    if(dims != fft_size)
    {
      fft_mem_offset=dims+2;
      fft_base::resize(dims,1);
#pragma omp critical
      {
        fftw_base<T>::clean();
        forward_plan=fftw_base<T>::create_plan(dims,1,in,mangle(in),uflag);
        backward_plan=fftw_base<T>::create_plan(dims,1,mangle(in),in,uflag);
      }
    }
  }

  inline void transform_forward(value_type* in, int howmany=1)
  {
    value_type* offset=in;
    for(int i=0; i<howmany; ++i,offset+=fft_mem_offset) 
      execute_dft(forward_plan,offset,mangle(offset));
  }

  inline void transform_backward(value_type* in, int howmany=1)
  {
    value_type* offset=in;
    for(int i=0; i<howmany; ++i,offset+=fft_mem_offset) 
      execute_dft(backward_plan,mangle(offset),offset);
  }
};

/** specialization for multiple-fft with FFTW on complex
 */
template<typename T>
struct FFT_DFT<T,std::complex<T>,true,FFTW_ENG> 
: public fft_base, public fftw_base<T>
{
  typedef std::complex<T> value_type;
  using fftw_base<T>::forward_plan;
  using fftw_base<T>::backward_plan;

  FFT_DFT(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    init(dims,in,howmany,uflag);
  }

  void init( int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    if(dims != fft_size)
    {
      fft_base::resize(dims,howmany);
#pragma omp critical
      {
        fftw_base<T>::clean();
        forward_plan = fftw_base<T>::create_plan(dims,howmany,mangle(in),mangle(in),FFTW_FORWARD,uflag);
        backward_plan = fftw_base<T>::create_plan(dims,howmany,mangle(in),mangle(in),FFTW_BACKWARD,uflag);
      }
    }
  }

  inline void transform_forward(value_type* in, int howmany=1)
  {
    fftw_base<T>::execute_dft(forward_plan,mangle(in),mangle(in));
  }

  inline void transform_backward(value_type* in, int howmany=1)
  {
    fftw_base<T>::execute_dft(backward_plan,mangle(in),mangle(in));
  }
};

/** specialization for multiple-fft with FFTW on double
 */
template<typename T>
struct FFT_DFT<T,T,true,FFTW_ENG> 
: public fft_base, public fftw_base<T>
{
  typedef T value_type;
  using fftw_base<T>::forward_plan;
  using fftw_base<T>::backward_plan;

  FFT_DFT() { }

  FFT_DFT(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    init(dims,in,howmany,uflag);
  }

  void init(int dims, value_type* in, int howmany=1, unsigned uflag=FFTW_MEASURE)
  {
    if(dims != fft_size)
    {
      fft_base::resize(dims,howmany);
#pragma omp critical
      {
        fftw_base<T>::clean();
        forward_plan = fftw_base<T>::create_plan(dims,howmany,in,mangle(in),uflag);
        backward_plan = fftw_base<T>::create_plan(dims,howmany,mangle(in),in,uflag);
      }
    }
  }

  inline void transform_forward(value_type* in, int howmany=1)
  {
    fftw_base<T>::execute_dft(forward_plan,in,mangle(in));
  }

  inline void transform_backward(value_type* in, int howmany=1)
  {
    fftw_base<T>::execute_dft(backward_plan,mangle(in),in);
  }
};

#endif
